{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install -qU memorizz yahooquery"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# ToolBox Functionality in the MemoRizz Library\n",
    "\n",
    "-------"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The Toolbox system in MemoRizz provides a powerful framework for registering, managing, and utilizing external functions as AI-callable tools. This system enables agents to seamlessly interact with external systems, APIs, and data sources through a standardized interface.\n",
    "\n",
    "## Core Concepts\n",
    "\n",
    "### Function Registration and Discovery\n",
    "- **Registration**: Convert Python functions into LLM-callable tools with semantic metadata\n",
    "- **Vector Embeddings**: Index tools for efficient semantic retrieval based on natural language queries\n",
    "- **Discoverability**: Find relevant tools based on contextual requirements without explicit references \n",
    "- **Persistence**: Store tool definitions across sessions and application restarts\n",
    "\n",
    "### Tool Management\n",
    "- **Centralized Repository**: Maintain all tools in a single, organized collection\n",
    "- **Access Control**: Configure private or global tool access patterns for agents\n",
    "- **Tool Metadata**: Enhance functions with rich descriptions and usage examples\n",
    "- **Versioning**: Update tool implementations while maintaining consistent identifiers\n",
    "\n",
    "### Agent Integration\n",
    "- **Dynamic Tool Selection**: Agents can discover and utilize the most relevant tools for each query\n",
    "- **Contextual Relevance**: Tools are matched to queries based on semantic similarity\n",
    "- **Automatic Documentation**: Generate parameter schemas and usage instructions for LLMs\n",
    "- **Error Handling**: Gracefully manage missing or broken tool implementations\n",
    "\n",
    "## Implementation Details\n",
    "\n",
    "Under the hood, the Toolbox system:\n",
    "1. Analyzes function signatures and docstrings to extract parameter information\n",
    "2. Generates embeddings that capture the semantic meaning of each tool\n",
    "3. Stores both the metadata and function references in the memory provider\n",
    "4. Provides retrieval mechanisms for both exact and similarity-based lookups\n",
    "\n",
    "This architecture enables a powerful \"tools as a service\" model where functions can be registered once and utilized by multiple agents throughout your application ecosystem.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Implementation & Usage Guidance\n",
    "\n",
    "\n",
    "To fully leverage MemoRizz's memory subsystem, always initialize a `MemoryProvider` before interacting with the library. In this workflow, the `MemoryProvider` acts as the authoritative store for all `ToolBox` data—handling both persistence and retrieval of function definitions, tool metadata, and embedding vectors throughout your application lifecycle. \n",
    "\n",
    "The ToolBox component relies on this persistent storage to enable seamless tool discovery, contextual matching, and function execution across different sessions and application instances.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import getpass\n",
    "import os\n",
    "\n",
    "# Function to securely get and set environment variables\n",
    "def set_env_securely(var_name, prompt):\n",
    "    value = getpass.getpass(prompt)\n",
    "    os.environ[var_name] = value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "set_env_securely(\"MONGODB_URI\", \"Enter your MongoDB URI: \")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "set_env_securely(\"OPENAI_API_KEY\", \"Enter your OpenAI API Key: \")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "from memorizz.memory_provider.mongodb.provider import MongoDBConfig, MongoDBProvider\n",
    "\n",
    "# Create a memory provider\n",
    "mongodb_config = MongoDBConfig(uri=os.environ[\"MONGODB_URI\"])\n",
    "memory_provider = MongoDBProvider(mongodb_config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 1: Creating Function-Based Tools\n",
    "\n",
    "In this initial step, we define specialized Python functions that will serve as the foundation for our toolbox. These functions encapsulate discrete capabilities that can be semantically discovered and executed by AI agents:\n",
    "\n",
    "1. The `get_stock_price()` function interfaces with the `yfinance` library to retrieve real-time financial market data. It accepts a stock symbol parameter (e.g., 'AAPL') and an optional currency code, returning formatted price information with proper numerical formatting.\n",
    "\n",
    "2. The `get_weather()` function leverages the Open-Meteo API to access meteorological data based on precise geolocation coordinates. It extracts the current temperature at the specified latitude and longitude, providing environmental context to agent interactions.\n",
    "\n",
    "These functions are designed with comprehensive type hints, descriptive docstrings, and robust error handling to ensure they can be properly indexed, discovered, and reliably executed when registered in the MemoRizz toolbox system.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "from functools import lru_cache\n",
    "from yahooquery import Ticker\n",
    "import time\n",
    "\n",
    "@lru_cache(maxsize=128)\n",
    "def _fetch_price(symbol: str) -> float:\n",
    "    \"\"\"\n",
    "    Internal helper to fetch the latest market price via yahooquery.\n",
    "    Caching helps avoid repeated hits for the same symbol.\n",
    "    \"\"\"\n",
    "    ticker = Ticker(symbol)\n",
    "    # This returns a dict keyed by symbol:\n",
    "    info = ticker.price or {}\n",
    "    # regularMarketPrice holds the current trading price\n",
    "    price = info.get(symbol.upper(), {}).get(\"regularMarketPrice\")\n",
    "    if price is None:\n",
    "        raise ValueError(f\"No price data for '{symbol}'\")\n",
    "    return price\n",
    "\n",
    "def get_stock_price(\n",
    "    symbol: str,\n",
    "    currency: str = \"USD\",\n",
    "    retry: int = 3,\n",
    "    backoff: float = 0.5\n",
    ") -> str:\n",
    "    \"\"\"\n",
    "    Get the current stock price for a given symbol using yahooquery,\n",
    "    with simple retry/backoff to handle occasional rate-limits.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    symbol : str\n",
    "        Stock ticker, e.g. \"AAPL\"\n",
    "    currency : str, optional\n",
    "        Currency code (Currently informational only; yahooquery returns native)\n",
    "    retry : int, optional\n",
    "        Number of retries on failure (default: 3)\n",
    "    backoff : float, optional\n",
    "        Backoff factor in seconds between retries (default: 0.5s)\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    str\n",
    "        e.g. \"The current price of AAPL is 172.34 USD.\"\n",
    "    \"\"\"\n",
    "    symbol = symbol.upper()\n",
    "    last_err = None\n",
    "    for attempt in range(1, retry + 1):\n",
    "        try:\n",
    "            price = _fetch_price(symbol)\n",
    "            return f\"The current price of {symbol} is {price:.2f} {currency.upper()}.\"\n",
    "        except Exception as e:\n",
    "            last_err = e\n",
    "            # simple backoff\n",
    "            time.sleep(backoff * attempt)\n",
    "    # if we get here, all retries failed\n",
    "    raise RuntimeError(f\"Failed to fetch price for '{symbol}' after {retry} attempts: {last_err}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(get_stock_price(\"AAPL\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "import requests\n",
    "\n",
    "def get_weather(latitude, longitude):\n",
    "    response = requests.get(f\"https://api.open-meteo.com/v1/forecast?latitude={latitude}&longitude={longitude}&current=temperature_2m,wind_speed_10m&hourly=temperature_2m,relative_humidity_2m,wind_speed_10m\")\n",
    "    data = response.json()\n",
    "    return data['current']['temperature_2m']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "latitude = 40.7128\n",
    "longitude = -74.0060\n",
    "weather = get_weather(latitude, longitude)\n",
    "print(weather)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Ensure values from the methods are printed in the cells above"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "### Step 2: Initializing the Toolbox and Registering Function Tools\n",
    "\n",
    "This phase focuses on transforming our raw Python functions into semantically discoverable AI tools through the MemoRizz orchestration layer:\n",
    "\n",
    "1. Instantiate a Toolbox instance that interfaces with our memory provider\n",
    "2. Register each function as a discoverable tool within the system\n",
    "3. Generate embedding vectors and metadata to enable contextual discovery\n",
    "4. Persist the registered tools in the memory provider for cross-session availability\n",
    "\n",
    "The registration process automatically analyzes function signatures, docstrings, and behavior patterns to create comprehensive tool schemas that language models can understand and invoke during agent interactions.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from memorizz import Toolbox\n",
    "\n",
    "toolbox = Toolbox(memory_provider=memory_provider)\n",
    "\n",
    "# Now the tools are registered in the memory provider within the toolbox\n",
    "toolbox.register_tool(get_weather)\n",
    "toolbox.register_tool(get_stock_price)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "### Step 3: Retrieving Tools from the Toolbox Repository\n",
    "\n",
    "The MemoRizz toolbox provides multiple access patterns for tool retrieval, accommodating different operational needs:\n",
    "\n",
    "1. **Deterministic Retrieval by Name**: The `get_tool_by_name()` method enables direct access to specific tools using their registered function names. This approach is ideal when you know exactly which tool you need.\n",
    "\n",
    "2. **Retrieval by Unique Identifier**: The `get_tool_by_id()` method allows precise tool access using the unique ID assigned during registration. This method is particularly useful in distributed systems or when tool names might overlap across different domains.\n",
    "\n",
    "3. **Semantic Discovery**: The `get_most_similar_tools()` method leverages vector embeddings to identify contextually relevant tools based on natural language queries. This powerful capability allows agents to discover appropriate tools without knowing their exact names.\n",
    "\n",
    "4. **Comprehensive Inventory**: The `list_tools()` method provides a complete catalog of all registered tools in the toolbox repository, enabling auditing, documentation generation, and management operations.\n",
    "\n",
    "Each retrieval method returns both the tool metadata (parameters, descriptions, etc.) and the callable Python function reference, ensuring seamless integration between the semantic discovery layer and executable functionality.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "specific_tool = toolbox.get_tool_by_name(\"get_stock_price\")\n",
    "specific_tool"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get similar tools\n",
    "similar_tools = toolbox.get_most_similar_tools(\"Get the stock price for Apple\", limit=1)\n",
    "similar_tools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "list_all_tools = toolbox.list_tools()\n",
    "list_all_tools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get a tool by id\n",
    "tool_id = similar_tools[0][\"_id\"]\n",
    "tool_retrieved_by_id = toolbox.get_tool_by_id(tool_id)\n",
    "tool_retrieved_by_id"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "### Creating a MemAgent with a Toolbox\n",
    "\n",
    "#### Design Considerations\n",
    "\n",
    "- A tool added via the `add_tool` method to the MemAgent will be located in the `tools` attribute of the agent itself. This approach ensures each agent maintains its own isolated tool registry.\n",
    "\n",
    "- A copy of the tool in the toolbox is made and stored within the MemAgent, preserving the tool's metadata, schema, and function reference. This design prevents cross-agent contamination when tool implementations are updated.\n",
    "\n",
    "- Methods such as `refresh_tools(tool_id)` allow you to update specific tools within an agent when their implementations change in the source toolbox. This ensures agents can receive tool updates without complete reconfiguration.\n",
    "\n",
    "- The `tool_access` attribute determines how tools are discovered during agent execution:\n",
    "  - When set to `\"private\"` (default), agents only access tools explicitly added to them\n",
    "  - When set to `\"global\"`, agents can dynamically discover any tool in the toolbox based on query relevance\n",
    "\n",
    "- Each agent maintains a private map of `tool_id → python_function` references in the `_tool_functions` attribute, ensuring function calls are properly directed to their implementations.\n",
    "\n",
    "- When an agent is persisted via `save()`, tool references are stored in the database, making them available after application restarts or when loading agents on different machines.\n",
    "\n",
    "- Tools with missing implementations are gracefully handled during agent execution with informative error messages rather than runtime failures.\n",
    "\n",
    "- You can inspect an agent's available tools using the native `list_tools()` method inherited from the toolbox interface.\n",
    "\n",
    "This architecture provides a clean separation between tool definition (in the toolbox) and tool availability (in the agent), enabling fine-grained control over which capabilities are exposed to different agents in your system.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a MemAgent and pass in the memory provider that the agent is stored in"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "from memorizz import MemAgent\n",
    "\n",
    "monday_agent = MemAgent(memory_provider=memory_provider)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save the agent to the memory provider\n",
    "monday_agent.save()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "The `add_tool` method can add tools to the Agent document in two ways:\n",
    "\n",
    "1. **Individual Tool Registration**: By providing a specific `tool_id` parameter, agents can selectively incorporate a single tool from the toolbox. This approach gives precise control over which capabilities are available to each agent.\n",
    "\n",
    "2. **Bulk Toolbox Integration**: By passing an entire `toolbox` instance, agents can automatically import all compatible tools in a single operation. This method is efficient when you want an agent to have access to a predefined set of related capabilities.\n",
    "\n",
    "In both cases, the system handles:\n",
    "- Resolving tool metadata from the memory provider\n",
    "- Connecting to the actual Python function implementations\n",
    "- Ensuring the tools are properly formatted for LLM function-calling\n",
    "- Persisting the tool references in the agent's configuration\n",
    "\n",
    "This flexibility allows developers to create specialized agents with targeted functionality or general-purpose agents with comprehensive tool access, all while maintaining a clean separation between tool definition and agent configuration.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "monday_agent.add_tool(toolbox=toolbox)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "monday_agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "monday_agent.run(\"Get me the stock price of Apple\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "monday_agent.run(\"Get me the weather in London\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 4: Updating Tools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "updated_tool = toolbox.update_tool_by_id(tool_id, {\"name\": \"get_stock_price_for_american_companies\"})\n",
    "updated_tool"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Confirm that the tool has been updated\n",
    "toolbox.get_tool_by_id(tool_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Delete a tool by name\n",
    "success = toolbox.delete_tool_by_name(\"get_stock_price\")\n",
    "success"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Delete a tool by id\n",
    "success_deleting_tool_by_id = toolbox.delete_tool_by_id(tool_id)\n",
    "success_deleting_tool_by_id\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Delete all tools within the toolbox\n",
    "success_deleting_all_tools = toolbox.delete_all()\n",
    "success_deleting_all_tools"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step  : Augmenting Tool Metadata"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can enhance the tool's searchability and retrieval accuracy by using an LLM to augment the tool description and generate synthetic queries. These enhancements are automatically added to the tool's metadata and used during embedding generation, significantly improving the tool selection process.\n",
    "\n",
    "To enable this feature, simply set `augment=True` when registering your tool using the `@toolbox.register_tool` decorator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "@toolbox.register_tool(augment=True)\n",
    "def updated_get_stock_price(ticker: str, currency: str = \"USD\") -> str:\n",
    "    \"\"\"\n",
    "    Get the current stock price for a specified ticker symbol.\n",
    "\n",
    "    Parameters:\n",
    "    -----------\n",
    "    ticker : str\n",
    "        The stock ticker symbol (e.g., \"AAPL\" for Apple).\n",
    "    currency : str, optional\n",
    "        The currency in which the price should be returned (default is USD).\n",
    "\n",
    "    Returns:\n",
    "    --------\n",
    "    str\n",
    "        A description of the current stock price.\n",
    "    \"\"\"\n",
    "    # Implementation here\n",
    "    pass"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "memorizz",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
